import json
import asyncio
from typing import Dict, Any, AsyncGenerator, Optional, Callable, Union, List

from fastapi import Request
from fastapi.responses import StreamingResponse
from httpx import Response
from starlette.background import BackgroundTask

from app.core.logger import logger

async def forward_stream_response(
    response: Response,
    transform_func: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None
) -> AsyncGenerator[str, None]:
    """
    将Dify API的SSE流式响应转发为FastAPI的SSE格式
    
    Args:
        response: 从Dify API获取的原始流式响应
        transform_func: 可选的转换函数，用于处理每个消息事件
        
    Yields:
        符合SSE格式的事件数据行
    """
    try:
        async for line in response.aiter_lines():
            line = line.strip()
            
            # 跳过空行
            if not line:
                continue
            
            # 解析SSE格式的行
            if line.startswith("data:"):
                event_data = line[5:].strip()
                if event_data:
                    try:
                        # 解析JSON数据
                        data = json.loads(event_data)
                        
                        # 应用转换函数（如果有）
                        if transform_func:
                            data = transform_func(data)
                            
                        # 返回SSE格式的事件
                        yield f"data: {json.dumps(data)}\n\n"
                        
                    except json.JSONDecodeError:
                        # 对于非JSON格式的行，直接转发
                        yield f"data: {event_data}\n\n"
            else:
                # 如果不是以data:开头，可能是其他SSE事件类型或自定义格式
                yield f"{line}\n\n"
                
    except Exception as e:
        logger.error(f"流式响应处理出错: {str(e)}")
        # 发送错误事件
        yield f"event: error\ndata: {json.dumps({'error': str(e)})}\n\n"
    finally:
        # 确保响应被关闭
        await response.aclose()

def create_streaming_response(
    response: Response,
    transform_func: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
    media_type: str = "text/event-stream"
) -> StreamingResponse:
    """
    创建流式响应
    
    Args:
        response: 从Dify API获取的原始流式响应
        transform_func: 可选的转换函数，用于处理每个消息事件
        media_type: 响应的媒体类型，默认为text/event-stream
        
    Returns:
        FastAPI的StreamingResponse对象
    """
    async def event_generator():
        async for chunk in forward_stream_response(response, transform_func):
            yield chunk
    
    # 创建后台任务来确保响应被关闭
    async def cleanup_response():
        await response.aclose()
    
    return StreamingResponse(
        event_generator(),
        media_type=media_type,
        background=BackgroundTask(cleanup_response)
    )

def combine_message_chunks(chunks: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    合并消息块
    
    用于将流式响应中的多个消息块合并为一个完整消息
    
    Args:
        chunks: 消息块列表
        
    Returns:
        合并后的完整消息
    """
    if not chunks:
        return {}
        
    # 使用最后一个消息块作为基础
    final_chunk = chunks[-1].copy()
    
    # 合并所有回答文本
    answers = [chunk.get("answer", "") for chunk in chunks if "answer" in chunk]
    if answers:
        final_chunk["answer"] = answers[-1]  # 使用最后一个块的answer
        
    # 合并所有代理思考
    agent_thoughts = []
    for chunk in chunks:
        if chunk.get("agent_thoughts"):
            for thought in chunk["agent_thoughts"]:
                # 检查是否已有相同代理的思考
                existing = next((t for t in agent_thoughts if t["agent_name"] == thought["agent_name"]), None)
                if existing:
                    # 更新现有思考
                    existing["thought"] = thought["thought"]
                else:
                    # 添加新思考
                    agent_thoughts.append(thought.copy())
    
    if agent_thoughts:
        final_chunk["agent_thoughts"] = agent_thoughts
        
    return final_chunk

async def buffer_stream_response(
    response: Response,
    transform_func: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None
) -> Dict[str, Any]:
    """
    缓冲流式响应并返回完整结果
    
    用于将流式响应转换为阻塞响应
    
    Args:
        response: 从Dify API获取的原始流式响应
        transform_func: 可选的转换函数，用于处理每个消息事件
        
    Returns:
        完整的响应结果
    """
    chunks = []
    
    try:
        async for line in response.aiter_lines():
            line = line.strip()
            
            # 跳过空行
            if not line:
                continue
            
            # 解析SSE格式的行
            if line.startswith("data:"):
                event_data = line[5:].strip()
                if event_data:
                    try:
                        # 解析JSON数据
                        data = json.loads(event_data)
                        
                        # 应用转换函数（如果有）
                        if transform_func:
                            data = transform_func(data)
                            
                        chunks.append(data)
                        
                    except json.JSONDecodeError:
                        # 对于非JSON数据，创建一个简单对象
                        chunks.append({"text": event_data, "raw": True})
    except Exception as e:
        logger.error(f"流式响应缓冲出错: {str(e)}")
        # 如果有错误，返回最后收集的数据或错误信息
        if not chunks:
            return {"error": str(e)}
    finally:
        # 确保响应被关闭
        await response.aclose()
    
    # 合并所有块
    if chunks:
        return combine_message_chunks(chunks)
    else:
        return {"error": "没有收到任何响应数据"} 